from typing import Annotated
from langchain_deepseek import ChatDeepSeek
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from dotenv import load_dotenv
import os
from langchain_tavily import TavilySearch


# 加载.env文件
os.environ["TAVILY_API_KEY"] = "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"

model = "deepseek-chat"
OPENAI_API_KEY = "sk-2f1767eebb5749adb135e7f93912a313"

tool = TavilySearch(max_results=2)
tools = [tool]
tool.invoke("LangGraph中的node是什么?")


llm = ChatDeepSeek(model=model, temperature=0, max_retries=2, api_key=OPENAI_API_KEY)
llm.bind_tools(tools)
class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}

graph_builder = StateGraph(State)

graph_builder.add_node("chatbot", chatbot)

graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break


# try:
#     display(Image(graph.get_graph().draw_mermaid_png()))
# except Exception:
#     # This requires some extra dependencies and is optional
#     pass



# messages = [
#     ("system", "你是一个翻译官，把用户的英语翻译成中文"),
#     ("human", "I love programming."),
# ]
#
# llm.invoke(messages)



